#!/usr/bin/env python

from collections import defaultdict
from subprocess import Popen, PIPE, STDOUT
from re import compile, VERBOSE
from os import remove, close
from string import split
from tempfile import mkstemp
from numpy import array

from ZenExceptions import ZenExternalException
from XYZRNWriter import XYZRNWriter
from ResidueProperties import AA_NAMES, NA_NAMES
from PDBParser import pdb2dict, match_coords
from Misc import deep_list, deep_tuple

def id_hetero(id_, location):
    """ Append hetero to residue name.
    """
    list_id = deep_list(id_)
    list_id[location][0] = 'H_' + list_id[location][0]
    return deep_tuple(list_id)

def get_full_id(atm):
    """ncat returns atomic contacts, translates them to atomic full id tuple, should
    comply with PDBParser
    """
    # we try here to mimimc zenPDBs full id.
    (model, chain_id, res_id, res_name, res_ic, at_name, at_elem, alt_loc) = atm
    model = int(model) - 1 # 0 conting
    if not chain_id:
        chain_id = " "      # re does not catch spaces
    if not alt_loc:
        alt_loc = " "       # re does not catch spaces

    res_id = int(res_id)    # can also be negative
    at_id = at_name.strip() #
    # we loose info about atom elements at_elem
    # could fail for CAlcium atoms in residues with CA atoms
    # but zenPDB will fail anyway if the is a C alpha (discard the calcium)

    # this will attach H_ to nucleotides... could be convenient
    # could introduce D for DNA P for PROTEIN and H for Hetero
    if res_name not in AA_NAMES and res_name not in NA_NAMES:
        res_name = '%s_%s' % ('H', res_name)
    # mimic id's
    at_long_id = (at_id, alt_loc)
    res_long_id = (res_name, res_id, res_ic)
    full_id = (model, chain_id, res_long_id, at_long_id)
    # we could autput all in a dict but the definite
    # outpu are atom-atom contacts
    return full_id

def guess_alt_loc(model, atom_full_id, atom_coords):
    """ Fixes issues with atom_ids missing alt_loc when a correct parsed model is there.
    """
    chn_id, res_id, at_id = atom_full_id[-3:]
    chn_id = (chn_id,)
    res_id = (res_id,)
    at_id = (at_id,)
    chain = model[chn_id]
    try:
        residue = chain[res_id]
    except KeyError:
        hetero_id = id_hetero(res_id, -1)
        residue = chain[hetero_id]
    try:
        atom = residue[at_id]
        return atom.get_full_id()
    except KeyError:
        for atom in residue:
            if not any(atom.get_coords() - atom_coords):
                return atom.get_full_id()
    return None

def call_areaimol(pdb_file1 ="", pdb_file2="", diffmode ="", mode ="",
                    smode ="", symmetry1 ="", symmetry2 ="",
                    trans ="", atom ="", exclude = "", matchup = "",
                    pntden = "", probe = ""):
    """ Calls AreaIMol from CCP4 to get the ASA values.
        This interface supports most of the options/functionality of
        AreaIMol. Note that all data is from the pseudo-pdb format
        output file rather than HTML summaries.
    """
    # other options need not be supported, accepts defaults if none given
    #
    if diffmode == 'compare' and not pdb_file2:
        print "Need two pdb_files"
        raise ZenExternalException
    dm_ = "DIFFMODE %s" % diffmode.upper()
    mo_ = "MODE %s" % mode.upper()
    sm_ = "SMODE %s" % smode.upper()
    sy_1 = "SYMMETRY %s" % symmetry1.upper()
    sy_2 = "SYMMETRY %s" % symmetry2.upper()
    tr_ = "TRANS %s" % trans.upper()
    at_ = "ATOM %s" % atom.upper()
    xc_ = "EXCLUDE %s" % exclude.upper()
    mu_ = "MATCHUP %s" % matchup.upper()
    pt_ = "PNTDEN %s" % pntden.upper()
    pb_ = "PROBE %s" % (probe.upper() or 1.4)
    print (dm_, mo_, sm_, sy_1, sy_2, tr_, at_, xc_, mu_, pt_, pb_, )
    # check if got option, and appends \n
    options = [opt + '\n' for opt in (dm_, mo_, sm_, sy_1, sy_2, tr_, at_, xc_, mu_, pt_, pb_, ) if len(split(opt)) is not 1]
    options.append('OUTPUT\n')
    print 'Running AreaIMol for %s with options:' % pdb_file1
    for option in options:
        print option,

    (fd, pdb_out) = mkstemp('.pdb')
    print pdb_out
    close(fd) # explicit close required
    # calling AreaIMol
    if diffmode.upper() == 'COMPARE': # need two input files
        process = Popen('areaimol xyzin %s xyzin2 %s xyzout %s' % (pdb_file1, pdb_file2, pdb_out), shell =True, stdin =PIPE, stdout =PIPE)
    else:
        process = Popen('areaimol xyzin %s xyzout %s' % (pdb_file1, pdb_out), shell =True, stdin =PIPE, stdout =PIPE)
    process.stdin.writelines(options)
    process.stdin.close()         # computation starts, eh.. should
    messages = process.stdout.read()         # argh! This output could also be parsed
    # TODO some errors could be cought here!
    tmp_file = open(pdb_out,'r')
    atom_lines = []
    for line in tmp_file.xreadlines():
        if not match_coords.match(line):
            continue
        fields = pdb2dict(line)
        data = fields['bfactor']

        coords = fields['coords']
        at_full_id = (0, fields['chain_id'], fields['res_long_id'], fields['at_long_id'])
        atom_lines.append((at_full_id, coords, data))
    tmp_file.close()
    remove(pdb_out)
    if not atom_lines:
        raise ZenExternalException
    return atom_lines

def call_msms(entities, probe_radius =1.4, surface ='tses', density =1.05, noh =False, radius_type ='united'):
    # density ... 1.05 well check 1G2A...
    def make_xyzrn(entities, radius_type):
        (fd_xyzrn, tmp_xyzrn) = mkstemp('xyzrn')
        close(fd_xyzrn)
        writer = XYZRNWriter()
        writer.set_structure(entities, radius_type)
        writer.save(tmp_xyzrn)
        return tmp_xyzrn

    def vert2array(filename):
        """ Read the vertex file into a numpy array.
        """
        try:
            handle = open(filename, "rb")
        except IOError:
            raise ZenExternalException, filename
        vertex_list = []
        for line in handle.xreadlines():
            elements = line.split()
            try:
                vertex = map(float, elements[0:3])
            except ValueError:
                continue
            vertex_list.append(vertex)
        handle.close()
        return array(vertex_list)

    tmp_xyzrn = make_xyzrn(entities, radius_type)
    (fd, tmp_of) = mkstemp('of')
    close(fd)

    #options
    if_ = '-if %s' % tmp_xyzrn
    of_ = '-of %s' % tmp_of
    af_ = '-af %s' % tmp_of # not a bug
    pr_ = '-probe_radius %s' % probe_radius
    su_ = '-surface %s' % surface
    de_ = '-density %s' % density
    if noh:
        nh_ = '-noh'
    else:
        nh_ = ''
    command = ' '.join(('msms', if_, of_, af_, pr_, su_, de_, nh_))
    print 'Running MSMS for %s with options:' % str(entities)
    print command[5:]

    process = Popen(command, shell=True, stdout=PIPE)
    output = process.stdout.readlines()
    # TODO - parse statistics

    dotvert = tmp_of + ".vert"
    dotface = tmp_of + ".face"
    dotarea = tmp_of + ".area"
    surface = vert2array(dotvert)

    # clean-up
    map(remove, [dotvert, dotface, dotarea, tmp_of, tmp_xyzrn])
    if not surface.any():
        raise ZenExternalException
    return surface

def call_stride(pdb_file):
    process = Popen('stride %s' % pdb_file, shell=True, stdin=PIPE, stdout=PIPE)
    output = {}
    for line in process.stdout.xreadlines():
        if line.startswith('ASG'):
            res_name = line[5:8] # we use 3
            chain_id = line[9]
            if chain_id == '-': # stride ' ' -> '-' rename
                chain_id = ' '
            try:
                res_id = float(line[10:15])
                res_ic = ' '
            except ValueError:
                res_id = float(line[10:14])
                res_ic = line[14]
            ss_code = line[24]
            phi = float(line[43:49])
            psi = float(line[53:59])
            asa = float(line[62:69])
            output[(chain_id, (res_name, res_id, res_ic))] = (ss_code, phi, psi, asa)
    if not output:
        # this propably means a Calpha only chain has been supplied
        raise ZenExternalException
    return output

def call_act(pdb_file ="", bdump ="", bmonitor ="", contact ="", hbond ="",
             sequ ="", short ="", symm = ""):
    """ Calls CCP4's ACT. ACT does not process the alt_loc identifier
        thus processing structure with alt_locs will result in ambiguous
        assignments. Support for insertion codes is also missing. ACT doesn't
        support PDBv3.0 files.
    """
    # other options need not be supported,
    # accepts defaults if none given.
    re_acterr = compile('ACT:')
    bd_ = "BDUMP %s" % bdump.upper()
    bm_ = "BMONITOR %s" % bmonitor.upper()
    co_ = "CONTACT %s" % contact.upper()
    sy_ = "SYMM %s" % symm.upper()
    sh_ = "SHORT %s" % short.upper()
    se_ = "SEQU %s" % sequ.upper()
    hb_ = "HBOND %s" % hbond.upper()
    # ACT
    print (bd_, bm_, co_, sy_, sh_, se_, se_, hb_)
    # check if got option, and appends \n
    options = [opt + '\n' for opt in (bd_, bm_, co_, sy_, sh_, se_, se_, hb_) if len(split(opt)) is not 1]
    options.append('END\n') # compulsory
    print 'Running ACT for %s with options:' % pdb_file
    for option in options:
        print option,
    # calling ACT
    process = Popen('act xyzin %s' % pdb_file, shell =True, stdin =PIPE, stdout =PIPE)
    process.stdin.writelines(options)
    process.stdin.close()
    lines = process.stdout.xreadlines()
    result = defaultdict(list)
    for line in lines:
        if line[18:22] == '....':
            if re_acterr.search(line):
                continue
            atoms = []
            for shift in (0, 21):
                res_id = line[shift+1:shift+6]
                chain_id = line[shift+7]
                res_name = line[shift+9:shift+12]
                at_name = line[shift+14:shift+18]
                atom = get_full_id((1, chain_id, res_id,
                                    res_name, ' ', at_name, ' ', ' '))
                atoms.append(atom)
            symmetry = line[39:53]
            distance = float(line[53:])
            result[atoms[0]].append((atoms[1], distance, symmetry))
    return dict(result)
